package easy.trace.plugin.crossthread.define;


import com.alibaba.ttl.TtlCallable;
import com.alibaba.ttl.TtlRunnable;
import easy.trace.agent.define.Instrumentation;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

import java.util.Objects;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @author easy
 */
public class CrossThreadInstrumentation implements Instrumentation {

    private static final String ENHANCE_TRANSFORM_CLASS = "java/util/concurrent/ThreadPoolExecutor";
    private static final String ENHANCE_CLASS = "java.util.concurrent.ThreadPoolExecutor";
    private static final String ENHANCE_SCHEDULED_CLASS = "java.util.concurrent.ScheduledThreadPoolExecutor";
    private static final String ENHANCE_TRANSFORM_SCHEDULED_CLASS = "java/util/concurrent/ScheduledThreadPoolExecutor";
    private static final String RUNNABLE_CLASS_NAME = "java.lang.Runnable";
    private static final String CALLABLE_CLASS_NAME = "java.util.concurrent.Callable";

    private AtomicBoolean isScheduledThread = new AtomicBoolean(false);

    @Override
    public byte[] interceptPoints() {
        final String enhanceClassName = isScheduledThread.get() ? ENHANCE_TRANSFORM_CLASS : ENHANCE_CLASS;
        System.out.println("---------------transform start:" + enhanceClassName + "---------------");
        try {
            ClassPool classPool = ClassPool.getDefault();
            CtClass ctClass;
            if (isScheduledThread.get()) {
                ctClass = classPool.getCtClass(ENHANCE_SCHEDULED_CLASS);
            } else {
                ctClass = classPool.getCtClass(ENHANCE_CLASS);
            }
            // 加载和执行生成的类
            for (CtMethod declaredMethod : ctClass.getDeclaredMethods()) {
                if (declaredMethod.getDeclaringClass() != ctClass) {
                    return new byte[0];
                }
                CtClass[] parameterTypes = declaredMethod.getParameterTypes();
                StringBuilder insertCode = new StringBuilder();
                for (int i = 0; i < parameterTypes.length; i++) {
                    CtClass paraType = parameterTypes[i];
                    if (RUNNABLE_CLASS_NAME.equals(paraType.getName())) {
                        String code = String.format("$%d = %s.get($%d, false, true);", i + 1, TtlRunnable.class.getName(), i + 1);
                        insertCode.append(code);
                    } else if (CALLABLE_CLASS_NAME.equals(paraType.getName())) {
                        String code = String.format("$%d = %s.get($%d, false, true);", i + 1, TtlCallable.class.getName(), i + 1);
                        insertCode.append(code);
                    }
                }
                if (insertCode.length() > 0) {
                    declaredMethod.insertBefore(insertCode.toString());
                }
            }
            System.out.println("---------------transform end:" + enhanceClassName + "---------------");
            return ctClass.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return new byte[0];
    }


    @Override
    public String getInterceptClass() {
        if (isScheduledThread.get()) {
            return ENHANCE_TRANSFORM_CLASS;
        } else {
            return ENHANCE_TRANSFORM_SCHEDULED_CLASS;
        }
    }

    @Override
    public boolean check(String className) {
        if (Objects.isNull(className) || className.length() == 0) {
            return false;
        }
        if (ENHANCE_TRANSFORM_CLASS.equalsIgnoreCase(className)) {
            isScheduledThread.set(false);
            return true;
        }
        if (ENHANCE_TRANSFORM_SCHEDULED_CLASS.equalsIgnoreCase(className)) {
            isScheduledThread.set(true);
            return true;
        }
        return false;
    }
}

